import torch # 导入 PyTorch 库，用于深度学习任务。
import gradio as gr  # 导入Gradio库，用于构建交互式界面。
import torch.nn.functional as F
from transformers import BertTokenizer, GPT2LMHeadModel # 从transformers库中导入BertTokenizer和GPT2LMHeadModel，用于自然语言处理和文本生成任务。

# Gitee AI 支持 transformers diffusers 等库，您可以直接使用 Gitee AI 模型库中的模型，前往 https://ai.gitee.com/models 查看更多模型。您也可以使用其他方式下载模型（例如 git 克隆、其他源下载）
tokenizer = BertTokenizer.from_pretrained("hf-models/gpt2-chinese-poem") # 用预训练的 BertTokenizer 加载中文诗歌生成模型的tokenizer。
model = GPT2LMHeadModel.from_pretrained("hf-models/gpt2-chinese-poem") # 使用预训练的GPT-2模型加载中文诗歌生成模型。
model.eval() # 将模型设置为评估模式，保证推理阶段的一致性、稳定性和效率。评估模式与训练模式相对应。

# 定义了一个函数，用于对模型输出的logits进行top-k和top-p过滤：
# 确保在生成文本时，只有概率最高的几个token被考虑，提高生成文本的质量和连贯性。
def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
    assert logits.dim() == 1
    top_k = min(top_k, logits.size(-1))
    if top_k > 0:
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value
    if top_p > 0.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[indices_to_remove] = filter_value
    return logits
# 定义了生成函数，接受一个输入文本，然后利用 AI 模型生成新文本。
def generate(input_text):
    input_ids = [tokenizer.cls_token_id]
    input_ids.extend( tokenizer.encode(input_text, add_special_tokens=False) )
    input_ids = torch.tensor( [input_ids] )
    generated = []
    for _ in range(100):
        output = model(input_ids)
        next_token_logits = output.logits[0, -1, :]
        next_token_logits[tokenizer.convert_tokens_to_ids('[UNK]')] = -float('Inf')
        filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=8, top_p=1)
        next_token = torch.multinomial( F.softmax(filtered_logits, dim=-1), num_samples=1 )
        if next_token == tokenizer.sep_token_id:
            break
        generated.append( next_token.item() )
        input_ids = torch.cat((input_ids, next_token.unsqueeze(0)), dim=1)
    return input_text + "".join( tokenizer.convert_ids_to_tokens(generated) )
examples = [["不堪翘首暮云中"], ["开源中国"], ["行到水穷处"], ["王师北定中原日"] ,["雪"], ["海上升明月"], ["十年磨一剑"]]

# 创建一个 Gradio 接口，用于接收用户输入，调用 generate 函数生成文本，并返回生成的文本。Gradio 默认启动到 7860 端口。
if __name__ == "__main__":
    gr.Interface(fn=generate, inputs="text", outputs="text",examples=examples).queue().launch()